#-*-coding:utf-8-*-#

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import container

os.chdir('PATH')

"""

Simulation results: Data_Fig4.zip

"""



dataframes = {}

nlist = [5,10,15,20,25,30,35,40,42,45,50]

for pnum in [2,3,4,5]:
    for i in nlist:
        dataframes[f'dr{i}_{pnum}'] = pd.read_csv('drought simulation AIDD t '+ str(i) +' number '+str(pnum)+'.csv',index_col=0)

nlist = [5,8,10,15,20,25,30,35,40,45,50]

for pnum in [2,3,4,5]:
    for i in nlist:
        dataframes[f'wf{i}_{pnum}'] = pd.read_csv('wildfire simulation AIDD t '+ str(i) +' number '+str(pnum)+'.csv',index_col=0)



def sens_test(df,simulation_num,leng):
    sdf = pd.DataFrame()
    sens = 0
    sign = 0
    for i in range(leng):
        if sens == 0 and 0 < df[str(simulation_num) + 't'].iloc[i]*100 and 0 < df[str(simulation_num) + 'b'].iloc[i]*100:
            sens = 1
        if sens == 1 and 0 < df[str(simulation_num) + 't'].iloc[i]*100 and 0 > df[str(simulation_num) + 'b'].iloc[i]*100:
            sign = i
            break
    if sens == 1:
        sdf.loc['vt','value'] = df[str(simulation_num) + 't'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vm','value'] = df[str(simulation_num) + 'm'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vb','value'] = df[str(simulation_num) + 'b'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['d','value'] = df[str(simulation_num) + 'm'].idxmax()
        if sign == 0:
            sdf.loc['sign','value'] = 194
        else:
            sdf.loc['sign','value'] = sign
        return sdf
    else:
        sdf.loc['vt','value'] = df[str(simulation_num) + 't'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vm','value'] = df[str(simulation_num) + 'm'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vb','value'] = df[str(simulation_num) + 'b'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['d','value'] = df[str(simulation_num) + 'm'].idxmax()
        sdf.loc['sign','value'] = -1
        return sdf

sendf = {}


dnlist = [5,10,15,20,25,30,35,40,42,45,50]
for pnum in [2,3,4,5]:
    for i in dnlist:
        if i != 8:
            sendf[f'wdrsdf{i}_{pnum}'] = sens_test(dataframes['dr'+str(i)+'_'+str(pnum)], 2, len(dataframes['dr'+str(i)+'_'+str(pnum)]))
        else:
            sendf[f'wdrsdf{i}_{pnum}'] = sens_test(dataframes['dr'+str(i)+'_'+str(pnum)], 2, len(dataframes['dr'+str(i)+'_'+str(pnum)]))

namelist = [f'wdrsdf{i}_{pnum}' for i in dnlist for pnum in [2,3,4,5]]

wdrsdf = pd.concat([sendf[f'wdrsdf{i}_{pnum}'] for i in dnlist for pnum in [2,3,4,5]],axis=1)
wdrsdf.columns = namelist
wdrsdf = wdrsdf.T


wnlist = [5,8,10,15,20,25,30,35,40,45,50]

for pnum in [2,3,4,5]:
    for i in wnlist:
        if i != 8:
            sendf[f'wwfsdf{i}_{pnum}'] = sens_test(dataframes['wf'+str(i)+'_'+str(pnum)]/100, 2, len(dataframes['wf'+str(i)+'_'+str(pnum)]))
        else:
            sendf[f'wwfsdf{i}_{pnum}'] = sens_test(dataframes['wf'+str(i)+'_'+str(pnum)]/100, 2, len(dataframes['wf'+str(i)+'_'+str(pnum)]))

namelist = [f'wwfsdf{i}_{pnum}' for i in wnlist for pnum in [2,3,4,5]]

wwfsdf = pd.concat([sendf[f'wwfsdf{i}_{pnum}'] for i in wnlist for pnum in [2,3,4,5]],axis=1)
wwfsdf.columns = namelist

wwfsdf = wwfsdf.T



"""


Extended data6a


"""



factors = ['2','3','4','5']
dlist1 = [str(x) for x in dnlist]

dl = dlist1 +dlist1 +dlist1 +dlist1

error1 = [wdrsdf['vm']-wdrsdf['vb'],wdrsdf['vt']-wdrsdf['vm']]

plt.figure(figsize=(15,5),dpi=400)
plt.errorbar(range(0,44,11),wdrsdf['vm'][:4],[error1[0][:4],error1[1][:4]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(1,44,11),wdrsdf['vm'][4:8],[error1[0][4:8],error1[1][4:8]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(2,44,11),wdrsdf['vm'][8:12],[error1[0][8:12],error1[1][8:12]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(3,44,11),wdrsdf['vm'][12:16],[error1[0][12:16],error1[1][12:16]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(4,44,11),wdrsdf['vm'][16:20],[error1[0][16:20],error1[1][16:20]],color='red',capsize = 5,marker='^',linestyle='none',)
plt.errorbar(range(5,44,11),wdrsdf['vm'][20:24],[error1[0][20:24],error1[1][20:24]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(6,44,11),wdrsdf['vm'][24:28],[error1[0][24:28],error1[1][24:28]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(7,44,11),wdrsdf['vm'][28:32],[error1[0][28:32],error1[1][28:32]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(8,44,11),wdrsdf['vm'][32:36],[error1[0][32:36],error1[1][32:36]],color='black',capsize = 5,marker='o',linestyle='none',label='Control (n=42)')
plt.errorbar(range(9,44,11),wdrsdf['vm'][36:40],[error1[0][36:40],error1[1][36:40]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(10,44,11),wdrsdf['vm'][40:],[error1[0][40:],error1[1][40:]],color='red',capsize = 5,marker='^',linestyle='none',label='Experiment')
plt.axvline(x=10.5,color='k',linewidth=1)
plt.axvline(x=21.5,color='k',linewidth=1)
plt.axvline(x=32.5,color='k',linewidth=1)
plt.axhline(y=0,color='black',linestyle='-',linewidth=2)
plt.axhline(y=10,color='black',linestyle='--',linewidth=1)
plt.axhline(y=20,color='black',linestyle='--',linewidth=1)
plt.axhline(y=30,color='black',linestyle='--',linewidth=1)
plt.axhline(y=40,color='black',linestyle='--',linewidth=1)

plt.xticks(ticks=range(0,44,1),labels=dl,fontsize=12)

plt.ylim([-10,45])
plt.xlim([-1,44])
ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
new_handles = []
for h in handles:
    #only need to edit the errorbar legend entries
    if isinstance(h, container.ErrorbarContainer):
        new_handles.append(h[0])
    else:
        new_handles.append(h)
ax.legend(new_handles, labels,prop = { "size": 12 }, loc ="upper left",framealpha=1,edgecolor='black')
plt.show()



"""


Extended data6a


"""


factors = ['2','3','4','5']
wflist1 = [str(x) for x in wnlist]
wfl = wflist1 + wflist1 + wflist1+ wflist1

error1 = [wwfsdf['vm']-wwfsdf['vb'],wwfsdf['vt']-wwfsdf['vm']]
fig, ax = plt.subplots(figsize=(15,5),dpi=400)
plt.errorbar(range(0,44,11),wwfsdf['vm'][:4],[error1[0][:4],error1[1][:4]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(1,44,11),wwfsdf['vm'][4:8],[error1[0][4:8],error1[1][4:8]],color='k',capsize = 5,marker='o',linestyle='none',label='Control (n=8)')
plt.errorbar(range(2,44,11),wwfsdf['vm'][8:12],[error1[0][8:12],error1[1][8:12]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(3,44,11),wwfsdf['vm'][12:16],[error1[0][12:16],error1[1][12:16]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(4,44,11),wwfsdf['vm'][16:20],[error1[0][16:20],error1[1][16:20]],color='red',capsize = 5,marker='^',linestyle='none',)
plt.errorbar(range(5,44,11),wwfsdf['vm'][20:24],[error1[0][20:24],error1[1][20:24]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(6,44,11),wwfsdf['vm'][24:28],[error1[0][24:28],error1[1][24:28]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(7,44,11),wwfsdf['vm'][28:32],[error1[0][28:32],error1[1][28:32]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(8,44,11),wwfsdf['vm'][32:36],[error1[0][32:36],error1[1][32:36]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(9,44,11),wwfsdf['vm'][36:40],[error1[0][36:40],error1[1][36:40]],color='red',capsize = 5,marker='^',linestyle='none')
plt.errorbar(range(10,44,11),wwfsdf['vm'][40:],[error1[0][40:],error1[1][40:]],color='red',capsize = 5,marker='^',linestyle='none',label='Experiment')
plt.axvline(x=10.5,color='k',linewidth=1)
plt.axvline(x=21.5,color='k',linewidth=1)
plt.axvline(x=32.5,color='k',linewidth=1)
plt.axhline(y=0,color='black',linestyle='-',linewidth=2)
plt.axhline(y=10,color='black',linestyle='--',linewidth=1)
plt.axhline(y=20,color='black',linestyle='--',linewidth=1)
plt.axhline(y=30,color='black',linestyle='--',linewidth=1)
plt.axhline(y=40,color='black',linestyle='--',linewidth=1)

plt.xticks(range(0,44,1), labels=wfl,fontsize=12)
plt.ylim([-10,45])
plt.xlim([-1,44])

ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
new_handles = []
for h in handles:
    #only need to edit the errorbar legend entries
    if isinstance(h, container.ErrorbarContainer):
        new_handles.append(h[0])
    else:
        new_handles.append(h)
ax.legend(new_handles, labels,prop = { "size": 12 }, loc ="upper left",framealpha=1,edgecolor='black',ncol=1)
plt.show()


